# -*- coding: utf-8 -*-
"""
Created on Wed Jun 26 12:10:15 2019

@author: CHD
"""
import os
from tensorflow.keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
'''
datagen = ImageDataGenerator(
        rotation_range=180,
        horizontal_flip = True,
        vertical_flip = True,
        fill_mode='nearest')
'''
datagen = ImageDataGenerator(
        rotation_range=180,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip = True,
        vertical_flip = True,
        fill_mode='nearest')


def Keras_image_argu(image_path,save_path):
    img = load_img(image_path)  # this is a PIL image
    x = img_to_array(img)  # this is a Numpy array with shape (3, 150, 150)
    x = x.reshape((1,) + x.shape)  # this is a Numpy array with shape (1, 3, 150, 150)
    # the .flow() command below generates batches of randomly transformed images
    # and saves the results to the `preview/` directory
    name = image_path.split('\\')[-1]
    name = name.split('.')[0]
    i = 0
    for batch in datagen.flow(x, batch_size=1,save_to_dir=save_path, save_prefix=name, save_format='jpg'):
        i += 1
        if i > 24:
            break  # otherwise the generator would loop indefinitely
       
def ImageArgu(path1, path2,save_path1,save_path2):
#读入文件列表，path代表2个类别的文件路径
    filelist1 = [os.path.join(path1, f) for f in os.listdir(path1)]
    filelist2 = [os.path.join(path2, f) for f in os.listdir(path2)]
    for img in filelist1:
        Keras_image_argu(img,save_path1)
    for img in filelist2:
        Keras_image_argu(img,save_path2)

path1 = '..\\trainset\\ClusterKDEpicture\\csv_dataset\\loc\\'
path2 = '..\\trainset\\NoClusterKDEpicture\\csv_dataset_no_cluster\\loc\\'     
save_path1 = '..\\trainset\\ArgumentImage\\singleOC\\'
save_path2 = '..\\trainset\\ArgumentImage\\singleField\\'
ImageArgu(path1, path2,save_path1,save_path2)